package com.lzugis;

import ucar.ma2.Array;
import ucar.ma2.ArrayFloat;
import ucar.ma2.DataType;
import ucar.ma2.Index;
import ucar.nc2.Dimension;
import ucar.nc2.NetcdfFile;
import ucar.nc2.NetcdfFileWriter;
import ucar.nc2.Variable;

import java.io.Serializable;
import java.util.ArrayList;

/**
 * @ClassName Grib2Test
 * @Author songyang
 * @Date 2022/5/3 22:51
 **/
public class Grib2Test {

    public static void main(String[] args) {
        try {
            NetcdfFile source = NetcdfFile.open("C:\\Users\\lzuni\\Desktop\\gfs.t12z.pgrb2.0p50.f000");
            Variable vu = source.findVariable("u-component_of_wind_height_above_ground");
            Variable vv = source.findVariable("v-component_of_wind_height_above_ground");
            Variable vlon = source.findVariable("lon");
            Variable vlat = source.findVariable("lat");

            String ncExtent = "0.0,359.5,0.5;90.0,-90.0,-0.5";
            float maxlon = 120.5f;
            float minlon = 118f;
            float maxlat = 37.5f;
            float minlat = 35.5f;
            GridIndexEntity startIndex = proj4j(minlon,minlat,ncExtent, -0.25);
            GridIndexEntity endIndex = proj4j(maxlon,maxlat,ncExtent, 0.25);

            String newNcFile = "C:\\Users\\lzuni\\Desktop\\aaa3.nc";
            NetcdfFileWriter aNew = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf3, newNcFile);
            Dimension time = aNew.addDimension(null, "time", 1);
            Dimension height = aNew.addDimension(null, "height", 7);
            Dimension lat = aNew.addDimension(null, "lat", startIndex.getY() - endIndex.getY());
            Dimension lon = aNew.addDimension(null, "lon", endIndex.getX() - startIndex.getX());
            ArrayList<Dimension> xythdims = new ArrayList<Dimension>();
            ArrayList<Dimension> xdim = new ArrayList<Dimension>();
            ArrayList<Dimension> ydim = new ArrayList<Dimension>();
            xdim.add(lon);
            ydim.add(lat);
            xythdims.add(time);
            xythdims.add(height);
            xythdims.add(lat);
            xythdims.add(lon);
            Variable newLon = aNew.addVariable(null, vlon.getShortName(), DataType.FLOAT, xdim);
            Variable newLat = aNew.addVariable(null, vlat.getShortName(), DataType.FLOAT, ydim);
            Variable newV =  aNew.addVariable(null, vv.getShortName(), DataType.FLOAT, xythdims);
            Variable newU =  aNew.addVariable(null, vu.getShortName(), DataType.FLOAT, xythdims);
            aNew.create();

            //写入经纬度
            ArrayFloat.D1 arrayFloatLat = (ArrayFloat.D1) vlat.read((endIndex.getY()) + ":" + (startIndex.getY() - 1));
            aNew.write(newLat, arrayFloatLat);
            ArrayFloat.D1 arrayFloatLon = (ArrayFloat.D1) vlon.read((startIndex.getX()) + ":" + (endIndex.getX() - 1));
            aNew.write(newLon, arrayFloatLon);

            //写入uv
            final String sectionSpec = "0:" + (vv.getShape()[0] - 1) + ":1,0:" + (vv.getShape()[1] - 1) + ":1," + (endIndex.getY()) + ":" + (startIndex.getY() - 1) + ":1,"
                    + (startIndex.getX()) + ":" + (endIndex.getX() - 1) + ":1";
            final Array readV = vv.read(sectionSpec);
            final Array readU = vu.read(sectionSpec);
            final int[] cacheShape = {1, vv.getShape()[1], startIndex.getY() - endIndex.getY(), endIndex.getX() - startIndex.getX()};
            final ArrayFloat.D4 arrayV = (ArrayFloat.D4) Array.factory(vv.getDataType(), cacheShape);
            final ArrayFloat.D4 arrayU = (ArrayFloat.D4) Array.factory(vu.getDataType(), cacheShape);
            final Index indexV = arrayV.getIndex();
            final Index indexU = arrayU.getIndex();
            final Index openIndexV = readV.getIndex();
            final Index openIndexU = readU.getIndex();

            for (int x = 0; x < cacheShape[1]; x++) {
                for (int j = 0; j < cacheShape[2] - 1; j++) {
                    for (int k = 0; k < cacheShape[3] - 1; k++) {
                        arrayV.setFloat(indexV.set(0, x, j, k), readV.getFloat(openIndexV.set(0, x, j, k)));
                        arrayU.setFloat(indexU.set(0, x, j, k), readU.getFloat(openIndexU.set(0, x, j, k)));
                    }
                }
            }
            aNew.write(newV, arrayV);
            aNew.write(newU, arrayU);
            aNew.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    public static GridIndexEntity proj4j(double lon, double lat, String ncExtent, double padding) {
        String[] extent = ncExtent.split(";");
        lon = lon + padding;
        lat = lat + padding;
        double minX = Double.parseDouble(extent[0].split(",")[0]);
        double minY = Double.parseDouble(extent[1].split(",")[0]);
        double spacingX = Double.parseDouble(extent[0].split(",")[2]);
        double spacingY = Double.parseDouble(extent[1].split(",")[2]);
        int xIndex = (int) ((lon - minX) / spacingX + 0.5);
        int yIndex = (int) ((lat - minY) / spacingY + 0.5);
        return new GridIndexEntity(xIndex, yIndex);
    }
    public static class GridIndexEntity implements Serializable {
        private static final long serialVersionUID = 1L;
        private int x;
        private int y;

        public int getX() {
            return x;
        }

        public void setX(int x) {
            this.x = x;
        }

        public int getY() {
            return y;
        }

        public void setY(int y) {
            this.y = y;
        }

        public GridIndexEntity(int x, int y) {
            this.x = x;
            this.y = y;
        }

        public GridIndexEntity() {
        }

        @Override
        public String toString() {
            return "GridIndexEntity{" +
                    "x=" + x +
                    ", y=" + y +
                    '}';
        }
    }
}
